from langchain.chains.retrieval_qa.base import RetrievalQA
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAI, OpenAIEmbeddings
from langchain_text_splitters import CharacterTextSplitter

# 加载文档
loader = TextLoader("d:\\data.txt", 'utf-8')
documents = loader.load()

# 分割文档
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
texts = text_splitter.split_documents(documents)
#
# 创建向量存储
embeddings = OpenAIEmbeddings(
    api_key="sk-VowKQBUMIkSND8WScNJtDLqf3FyqWHQ43LMVUXH1m6GZaopA",
    base_url="https://ai.nengyongai.cn/v1",
    model="text-embedding-3-small"
)

texts = ["我是王振军", "是一名软件工程师", "精通java python语言，工作认真踏实"]
vector_store = FAISS.from_texts(texts, embeddings)

llm = OpenAI(
    api_key="sk-VowKQBUMIkSND8WScNJtDLqf3FyqWHQ43LMVUXH1m6GZaopA",
    base_url="https://ai.nengyongai.cn/v1"
)
# 3. 构建检索式问答链
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=vector_store.as_retriever()
)
# 4. 提问
result = qa_chain.invoke("王振军是谁？")
print(result)
